import sys, os
from ros import rosbag
import roslib
import rospy
roslib.load_manifest('sensor_msgs')
from sensor_msgs.msg import Image,Imu
from geometry_msgs.msg import Vector3

from PIL import ImageFile
from PIL import Image as ImagePIL

def readImg(folder):
    print("Searching directory %s" % folder)
    all = []
    files = os.listdir(folder)
    for f in sorted(files):
        if os.path.splitext(f)[1] in ['.png', '.jpg']:
            all.append(os.path.join(folder,f))
    # print(len(all))
    return all

def readTimestamp(file):
    print("timestampfile %s" % file)
    f = open(file,"r")
    data = f.readlines()
    f.close()
    vtimestamp=[]
    for i in range(len(data)):
        tmp = data[i].split(' ')
        vtimestamp.append(tmp[0])
    print(len(vtimestamp))
    return vtimestamp


def createBag(args):
    print("bag will be created")
    vimgs = readImg(args[0])
    imagetimestamps = readTimestamp(args[1])
    if not os.path.exists(args[2]):
        os.system(r'touch %s' % args[2])
    bag = rosbag.Bag(args[2],'w')

    try:
        for i in range(len(vimgs)):
            print("adding image %s" % vimgs[i])
            fp = open(vimgs[i], "r")
            p = ImageFile.Parser()

            imgpil = ImagePIL.open(vimgs[i])
            width, height = imgpil.size

            while 1:
                s = fp.read(1024)
                if not s:
                    break
                p.feed(s)

            im = p.close()

            stamp = rospy.rostime.Time.from_sec(float(imagetimestamps[i])/1e3)

            Img = Image()
            Img.header.stamp = stamp
            Img.height = height
            Img.header.frame_id = "cam0"

            Img.encoding = "mono8"
            Img_data = [pix for pixdata in [im.getdata()] for pix in pixdata]
            Img.step = Img.width

            Img.data = Img_data
            bag.write('/cam0/image_raw',Img, stamp)
            # bag.write('/cam1/image_raw',Img, stamp)
    finally:
        bag.close()

if __name__ == "__main__":
    print(sys.argv)
    createBag(sys.argv[1:])